# -*- coding: utf-8 -*-
"""
@Time ： 2020-01-06 10:46
@Auth ： chenzj85
@Description：

"""
import shutil

import imageio
import logging
from algo.Algo_interface import Algo_interface
import json
import sys
import os
import glob
from imageai.Detection.Custom import DetectionModelTrainer
from imageai.Detection.Custom import CustomObjectDetection
import numpy as np
import numpy as np
import cv2
import random
import pandas as pd
from tensorflow import keras

class ObjectDetection(Algo_interface):
    def __init__(self, model_type, model_name, model_params):
        self.task_type = model_type
        self.model_name = model_name
        self.model_params = model_params
        self.model = dict()
        self.build_model()
        # return self.model


    def set_model(self, model):
        self.model = model
        return 1

    def get_model(self):
        self.model.pop('train')
        self.model.pop('model')
        return self.model

    def build_model(self):

        if self.model_name == 'yolo':
            self.model['train'] = DetectionModelTrainer()
            self.model['train'].setModelTypeAsYOLOv3()

    def train(self,data):
        def new_report(test_report):
            lists = os.listdir(test_report)  # 列出目录的下所有文件和文件夹保存到lists
            lists.sort(key=lambda fn: os.path.getmtime(test_report + "/" + fn))  # 按时间排序
            file_new = os.path.join(test_report, lists[-1])  # 获取最新的文件保存到file_new
            return file_new

        try:
            xml_train_path,lable_map = data  # 在api.py是tuple进来的两个Series
            self.model['train'].setDataDirectory(data_directory=xml_train_path)
            lable_list = list(lable_map['label'].keys())
            self.model['train'].setTrainConfig(object_names_array=lable_list, batch_size= self.model_params['batch_size'], num_experiments=self.model_params['num_experiments'])
            # In the above,when training for detecting multiple objects,
            # set object_names_array=["object1", "object2", "object3",..."objectz"]
            self.model['train'].trainModel()
            self.model['json'] = os.path.abspath(xml_train_path + '/json/detection_config.json')
            self.model['model_path'] = os.path.abspath(new_report(xml_train_path + '/models'))
            ##训练后的模型加载
            self.model['model'] = CustomObjectDetection()
            self.model['model'].setModelTypeAsYOLOv3()
            self.model['model'].setJsonPath(self.model['json'])
            self.model['model'].setModelPath(self.model['model_path'])
            self.model['model'].loadModel()

        except Exception as e:
            print(e)

        return 1

    def predict(self,data):   #data为图片路径list
        data,type = data
        if type == 'evaluate':
            def box_detect(image_path):
                detections = self.model['model'].detectObjectsFromImage(input_image=image_path,output_image_path="/opt/AiStudio/ai-platform-training-master/image/result.jpg", minimum_percentage_probability=30)
                predict_box = []
                for eachObject in detections:
                    predict_box.append(eachObject["box_points"])
                return predict_box
            #返回data用于计算出训练端准确率
            data['box_predict'] = list(map(lambda x: box_detect(x), data['image_path']))
            return data


        elif type == 'predict':
            #返回result的dataframe为预测结果可直接写进csv
            result = []
            image_path = list(data['image_path'])
            for image in image_path:
                detections = self.model['model'].detectObjectsFromImage(input_image=image,output_image_path="result.jpg", minimum_percentage_probability=30)
                if len(detections) == 0:
                    result.append([os.path.basename(image), None, None])
                else:
                    for eachObject in detections:
                        result.append([os.path.basename(image),eachObject["name"],eachObject["box_points"]])

            return pd.DataFrame(result,columns=['image','label','box_points'])


